const dotenv = require('dotenv');
const path = require('path');
const fs = require('fs');

// 加载环境变量
const envPath = path.resolve(__dirname, '.env');
dotenv.config({ path: envPath });

const { HNSWLib } = require("@langchain/community/vectorstores/hnswlib");
const { Embeddings } = require("@langchain/core/embeddings");

// 自定义嵌入类，返回固定向量用于测试
class TestEmbeddings extends Embeddings {
    async embedDocuments(texts) {
        return texts.map(() => Array(1536).fill(0.1));
    }

    async embedQuery(text) {
        return Array(1536).fill(0.1);
    }
}

// 修改向量存储路径为正确的路径
const VECTOR_STORE_PATH = path.join(__dirname, '../vector_store');
const embeddings = new TestEmbeddings();

async function testHNSW() {
    try {
        console.log("开始测试HNSWLib...");
        
        // 确保目录存在
        if (!fs.existsSync(VECTOR_STORE_PATH)) {
            fs.mkdirSync(VECTOR_STORE_PATH, { recursive: true });
            console.log("创建向量存储目录");
        }
        
        // 创建测试文档
        const documents = [
            { pageContent: "这是第一个测试文档", metadata: { id: "1" } },
            { pageContent: "这是第二个测试文档", metadata: { id: "2" } }
        ];
        
        console.log("创建向量存储...");
        const vectorStore = await HNSWLib.fromDocuments(documents, embeddings);
        
        console.log("保存向量存储...");
        await vectorStore.save(VECTOR_STORE_PATH);
        
        // 检查文件是否创建成功
        console.log("检查生成的文件:");
        const files = fs.readdirSync(VECTOR_STORE_PATH);
        for (const file of files) {
            const filePath = path.join(VECTOR_STORE_PATH, file);
            const stats = fs.statSync(filePath);
            console.log(`${file}: ${stats.size} 字节`);
        }
        
        console.log("加载向量存储...");
        const loadedVectorStore = await HNSWLib.load(VECTOR_STORE_PATH, embeddings);
        
        console.log("执行相似度搜索...");
        const results = await loadedVectorStore.similaritySearch("测试查询", 1);
        
        console.log("搜索结果:", results);
        console.log("测试完成！");
    } catch (error) {
        console.error("测试失败:", error);
    }
}

testHNSW();